import json
import os
import time
from glob import glob

import cv2
import mmcv
# using a pre-trained detector. 预训练配置
from mmcv import Config
from mmcv.runner import load_checkpoint

from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector, inference_detector, show_result_pyplot

from mmdet.apis import init_detector, inference_detector, show_result_pyplot


def init():

    # root_path=r"/project/train/src_repo/mmdetection/"
    root_path = r"/home/deepin/Documents/ji_pingtai/mmdetection/"

    config_file = root_path + 'work_dir/cocoDataset_cfgformat.py'

    checkpoint_file = root_path + 'work_dir/latest.pth'

    device = 'cuda:0'

    # 初始化检测器--构建模型
    model = init_detector(config_file, checkpoint_file, device=device)

    # # 查看 faster RCNN模型结构：
    # for name, module in model.named_children():
    #     print(name)
    #     [print(F'    {n}') for n, _ in module.named_children()]
    #
    # # 推理 显示
    # # img = mmcv.imread(root_path+'data/coco/val2017/ZDSmask20220829_V8_train_office_2_003505.jpg')
    # img = mmcv.imread(root_path + 'data/coco/val2017/2.jpg')
    #
    # result = inference_detector(model, img)  # ------- 训练完成 加载 ，推理
    #
    # print(result)
    #
    # # 先挑出来，conf_thres>= 0.3 的
    # # Load checkpoint
    # checkpoint = load_checkpoint(model, checkpoint_file, map_location=device)
    #
    # import numpy as np
    # list_bboxs = []
    # conf_thres = 0.5  # --------------------------------置信度
    # for j in range(len(result)):
    #
    #     for i, x in np.ndenumerate(result[j]):
    #
    #         # print(i, x) # (1, 4) 0.25840828 ---每个 坐标的 值，及 坐标
    #
    #         if i[1] == 4 and x >= conf_thres:
    #             cl_numb = i[0]  # 行号
    #             box = result[j][cl_numb]  # 找到 大于 置信度 那一行
    #
    #             left = box[0]
    #             top = box[1]
    #
    #             width = box[2] - box[0]
    #             height = box[3] - box[1]
    #             confidence = float(box[4])
    #             name = checkpoint['meta']['CLASSES'][j]
    #
    #             list_bboxs.append((left, top, width, height, name, confidence))
    #
    # print(list_bboxs)
    #
    # # show_result_pyplot(model, img, result, score_thr=0.3, out_file=root_path + 'data/result.jpg')  # 显示


    return model




def process_image(model, input_image=None, args=None, **kwargs):

    # 类别
    class_names = ['front_wear', 'front_no_wear', 'front_under_nose_wear', 'front_under_mouth_wear',
                    'front_unknown', 'side_wear', 'side_no_wear', 'side_under_nose_wear',
                    'side_under_mouth_wear', 'side_unknown', 'side_back_head_wear', 'side_back_head_no_wear',
                    'back_head', 'mask_front_wear', 'mask_front_under_nose_wear',
                    'mask_front_under_mouth_wear', 'mask_side_wear', 'mask_side_under_nose_wear',
                    'mask_side_under_mouth_wear', 'strap']

    # 推理 显示
    # img = mmcv.imread(root_path+'data/coco/val2017/ZDSmask20220829_V8_train_office_2_003505.jpg')
    # img =cv2.imread(image_name) # mmcv.imread(input_image) # 输入 mmcv.imread(root_path + 'data/coco/val2017/2.jpg')
    img=input_image

    result = inference_detector(model, img)  # ------- 训练完成 加载 ，推理

    # print(result)

    # 先挑出来，conf_thres>= 0.3 的

    import numpy as np
    list_bboxs = []
    conf_thres = 0.5  # --------------------------------置信度
    for j in range(len(result)):

        for i, x in np.ndenumerate(result[j]):

            # print(i, x) # (1, 4) 0.25840828 ---每个 坐标的 值，及 坐标

            if i[1] == 4 and x >= conf_thres:
                cl_numb = i[0]  # 行号
                box = result[j][cl_numb].tolist()  # 找到 大于 置信度 那一行, .tolist()--将numpy.array转换成list

                left = box[0]
                top = box[1]

                width = box[2] - box[0]
                height = box[3] - box[1]
                confidence =box[4]
                name = class_names[j]

                list_bboxs.append((left, top, width, height, name, confidence))

    # print(list_bboxs)

    fake_result = {}

    fake_result["algorithm_data"] = {  # 也就是说，这个这里 只是带上的，不要管
        "is_alert": False,
        "target_count": 0,
        "target_info": []
    }

    fake_result["model_data"] = {"objects": []}

    # Process detections
    cnt = 0
    # 没带口罩的 标签
    no_waers = ['front_no_wear', 'front_under_nose_wear', 'front_under_mouth_wear',
                'side_no_wear', 'side_under_nose_wear', 'side_under_mouth_wear']

    for box in list_bboxs:
        # box = list_bboxs[i]
        left = box[0]
        top = box[1]

        width = box[2]
        height = box[3]
        confidence = str(box[5])
        name = box[4]

        if name in no_waers:  # 没有 带 口罩
            # print(name)
            cnt += 1
            fake_result["model_data"]['objects'].append({  # 这里 才是，真正要添加的
                "x": left,
                "y": top,
                "height": height,
                "width": width,
                "confidence": confidence,
                "name": name
            })
            fake_result["algorithm_data"]["is_alert"] = True
            fake_result["algorithm_data"]["target_count"] = cnt

            fake_result["algorithm_data"]["target_info"].append({
                "x": left,
                "y": top,
                "height": height,
                "width": width,
                "confidence": confidence,
                "name": name
            })

        else:  # 带了 口罩
            # print(name)
            fake_result["model_data"]['objects'].append({  # 这里 才是，真正要添加的
                "x": left,
                "y": top,
                "height": height,
                "width": width,
                "confidence": confidence,
                "name": name
            })

            fake_result["algorithm_data"] = {  # 也就是说，这个这里 只是带上的，不要管
                "is_alert": False,
                "target_count": 0,
                "target_info": []
            }

    return json.dumps(fake_result, indent=4)


if __name__ == '__main__':

    # PLUGIN_LIBRARY = "weights/libmyplugins.so"
    # ctypes.CDLL(PLUGIN_LIBRARY)

    # image_names = glob(r'/project/ev_sdk/data/1.jpg')
    image_names = glob(r'/home/deepin/Documents/ji_pingtai/mmdetection/data/coco/val2017/2.jpg')

    # modelpath_car = r'/home/deepin/Documents/ji_pingtai/windows/yolov5/runs/train/exp/weights/best.onnx'
    # image_names = glob(r'/home/deepin/Documents/ji_pingtai/windows/yolov5/data/images/bus.jpg')

    predictor = init() # 初始化，加载 模型

    s = 0
    count = 0
    for index in range(2):
        for image_name in image_names:
            # print('image_path:', os.path.join(image_dir, image_name))
            img =  cv2.imread(image_name)

            start = time.time()

            res = process_image(predictor, img)

            print(res)

            end = time.time()

            inf_end = end - start
            s += end - start

            fps = 1 / inf_end

            fps_label = "FPS: %.2f " % fps
            fps_label = fps_label + " %.2f" % inf_end + "s " + f'({((float(inf_end) * 1000.0)):.1f}ms) Inference'
            print(fps_label)
            print('--------------------------------------------------------:', index)













